// Copyright (c) 1996-2002 John Lyon-Smith. All rights reserved.

#pragma once

#include <wxtl\Base.h>
#include <wxtl\SockAddrIn.h>
#include <wxtl\Event.h>

#include <MSWSock.h>

#pragma comment(lib, "ws2_32")

#define THROW_BAD_SOCKET_RET(x)	{ if ((x) == SOCKET_ERROR && WSAGetLastError() != WSA_IO_PENDING) throw _win32_error(WSAGetLastError()); }

namespace wxtl
{
	class Socket
	{
	public:
		
		bool TransmitFile(File& file, DWORD fileSize, DWORD numberOfBytesPerSend, LPOVERLAPPED overlapped, LPTRANSMIT_FILE_BUFFERS transmitBuffers, DWORD flags);

		Socket() : m_socket(INVALID_SOCKET)
		{ 
			clearAddrs();
		};
		Socket(Socket& h) throw() :
			m_socket(INVALID_SOCKET)
		{
			wxtl::SockAddrIn remoteAddr = h.m_remoteAddr;
			wxtl::SockAddrIn localAddr = h.m_localAddr;
			
			attach(h.detach());

			m_remoteAddr = remoteAddr;
			m_localAddr = localAddr;
		}

		Socket(SOCKET s) throw() :
			m_socket(s)
		{
			clearAddrs();
		}

		virtual ~Socket() 
		{ 
			close(); 
		};

		Socket& operator=(Socket& rhs) throw()
		{
			if (this != &rhs)
			{
				// Save the rhs addresses
				m_remoteAddr = rhs.m_remoteAddr;
				m_localAddr = rhs.m_remoteAddr;
				
				attach(rhs.detach());
			}

			return *this;
		}

		void attach(SOCKET socket) throw()
		{
			close();
			
			m_socket = socket;
		}

		SOCKET detach() throw()
		{
			SOCKET sock = m_socket;

			// Do the same as close without disconnecting
			m_socket = INVALID_SOCKET;
			m_evtNetwork.close();
			clearAddrs();

			return sock;
		}

		bool isValid() const throw()
		{ 
			return INVALID_SOCKET != m_socket; 
		}

		void close() throw()
		{
			if (isValid())
			{
				::shutdown(m_socket, SD_SEND);
				::closesocket(m_socket);
			}

			m_socket = INVALID_SOCKET;
			m_evtNetwork.close();
			clearAddrs();
		}

		SOCKET get() const throw()
		{ 
			return m_socket; 
		}

		enum
		{
			Stream = SOCK_STREAM,
			DataGram = SOCK_DGRAM,
		};

		void create(int nSocketType = Stream, bool bOverlapped = false) throw(...)
		{
			m_socket = WSASocket(AF_INET, nSocketType, IPPROTO_TCP, NULL, 0, 
				bOverlapped ? WSA_FLAG_OVERLAPPED : 0);

			if (INVALID_SOCKET == m_socket)
				throw _win32_error(WSAGetLastError());
		}

		void bind(SockAddrIn &addr) throw(...)
		{
			_ASSERTE(isValid());

			m_localAddr = addr;

			THROW_BAD_SOCKET_RET(::bind(m_socket, (LPSOCKADDR) &m_localAddr, sizeof(m_localAddr)));
		}

		void bind(UINT nSocketPort = 0, const TCHAR* pszSockAddr = NULL) throw(...)
		{
			m_localAddr.set(nSocketPort, pszSockAddr);

			THROW_BAD_SOCKET_RET(::bind(m_socket, (LPSOCKADDR) &m_localAddr, sizeof(m_localAddr)));
		}

		void eventSelect(LONG lBitMask) throw(...)
		{
			_ASSERTE(isValid());

			// Create the event if not yet created
			if (!m_evtNetwork.isValid())
				m_evtNetwork.create();

			THROW_BAD_SOCKET_RET(WSAEventSelect(m_socket, m_evtNetwork.get(), lBitMask));
		}

		void setNagle(bool bNagle) throw(...)
		{
			long lData = !bNagle;

			THROW_BAD_SOCKET_RET(::setsockopt(m_socket, IPPROTO_TCP, TCP_NODELAY, (char*)&lData, sizeof(lData)));
		}

		void listen(UINT nBackLog = SOMAXCONN) throw(...)
		{
			_ASSERTE(isValid());

			THROW_BAD_SOCKET_RET(::listen(m_socket, nBackLog));
		}

		void accept(Socket* pConnect) throw(...)
		{
			if(pConnect == NULL)
			{
				abort(); // TODO Log here
			}
			
			SOCKET sockAccept;
			SockAddrIn remoteAddr;
			int nAddrLen;

			// Make sure we have an empty socket object for the connection
			_ASSERTE(!pConnect->isValid());

			// accept the connection
			nAddrLen = sizeof(remoteAddr);
			sockAccept = ::accept(m_socket, remoteAddr.ptr(), &nAddrLen);
			THROW_BAD_SOCKET_RET(sockAccept);

			// Set up the connected socket object
			pConnect->m_socket = sockAccept;
			pConnect->m_remoteAddr = remoteAddr;
		}

		void connect(UINT port, const TCHAR* pSockAddr) throw(...)
		{
			_ASSERTE(isValid());
			_ASSERTE(0 != port);
			_ASSERTE(NULL != pSockAddr);

			SockAddrIn addr;

			addr.set(port, pSockAddr);

			connect(addr);
		}

		void connect(UINT port, const CHAR* pSockAddr, LPOVERLAPPED pOverlap) throw(...);

		void connect(SockAddrIn addr) throw(...)
		{
			THROW_BAD_SOCKET_RET(::connect(m_socket, addr.ptr(), sizeof(SockAddrIn)));

			m_remoteAddr = addr;
		}

		bool recv(void* pBuff, size_t bufLen, size_t* pTransferred = NULL,
			DWORD *pdwFlags = NULL, LPOVERLAPPED pOverlap = NULL) throw(...)
		{
			WSABUF wsab;

			wsab.buf = (char*)pBuff;
			wsab.len = (DWORD)bufLen;

			return recvScatter(&wsab, 1, pTransferred, pdwFlags, pOverlap);
		}

		bool IsConnected()
		{
			DWORD time;
			int optlen = sizeof(DWORD);
			if(getsockopt(m_socket, SOL_SOCKET, SO_CONNECT_TIME,(char*)&time, &optlen) 
								==SOCKET_ERROR)
			{
				//int error = WSAGetLastError();
				return false;
			}
			else
			{
				return true;
			}
		}

		bool send(const void* pBuff, size_t bufLen, size_t* pTransferred = NULL, 
			DWORD dwFlags = 0, LPOVERLAPPED pOverlap = NULL) throw(...)
		{
			WSABUF wsab;

			wsab.buf = (char*)const_cast<void*>(pBuff);
			wsab.len = (DWORD)bufLen;

			return sendGather(&wsab, 1, pTransferred, dwFlags, pOverlap);
		}

		bool recvScatter(LPWSABUF pBuf, size_t bufLen, 
			size_t* pTransfer = 0, DWORD* pdwFlags = 0, LPOVERLAPPED pOverlap = NULL) throw(...)
		{
			_ASSERTE(isValid());
			
			DWORD dwTransfer;
			DWORD dwFlags = pdwFlags != NULL ? *pdwFlags : 0;
			
			DWORD ret = WSARecv(m_socket, pBuf, (DWORD)bufLen, 
				&dwTransfer, &dwFlags, pOverlap, NULL);
			THROW_BAD_SOCKET_RET(ret);

			if (pTransfer)
				*pTransfer = (size_t)dwTransfer;

			if (pdwFlags)
				*pdwFlags = dwFlags;

			return ret != 0;
		}
		
		bool sendGather(LPWSABUF pBuf, size_t bufLen, 
			size_t* pTransfer = 0, DWORD dwFlags = 0, LPOVERLAPPED pOverlap = NULL) throw(...)
		{
			_ASSERTE(isValid());
			
			DWORD dwTransfer;
			
			DWORD ret = WSASend(m_socket, pBuf, (DWORD)bufLen, 
				&dwTransfer, dwFlags, pOverlap, NULL);
			THROW_BAD_SOCKET_RET(ret);

			if (pTransfer)
				*pTransfer = (size_t)dwTransfer;

			return ret != 0;
		}
		
		void recvFromScatter(LPWSABUF pBuf, DWORD dwBufLen, 
			PDWORD pdwTransfer, SockAddrIn& fromAddr, PDWORD pdwFlags = 0, LPOVERLAPPED pOverlap = NULL) throw(...)
		{
			int addrLen = sizeof(fromAddr);

			THROW_BAD_SOCKET_RET(WSARecvFrom(
				m_socket, pBuf, dwBufLen, pdwTransfer, pdwFlags, 
				fromAddr.ptr(), &addrLen, pOverlap, NULL));
		}

		void sendToGather(LPWSABUF pBuf, DWORD dwBufLen, 
			PDWORD pdwTransfer, SockAddrIn& toAddr, DWORD dwFlags = 0, LPOVERLAPPED pOverlap = NULL) throw(...)
		{
			THROW_BAD_SOCKET_RET(WSASendTo(
				m_socket, pBuf, dwBufLen, pdwTransfer, dwFlags, 
				toAddr.ptr(), sizeof(toAddr), pOverlap, NULL));
		}

		const SockAddrIn& getLocalAddr()
		{ 
			return m_localAddr;
		}
		const SockAddrIn& getRemoteAddr()
		{
			return m_remoteAddr;
		}

		


	// Static methods
	
		static void startUp(BYTE major = 2, BYTE minor = 0, WSADATA *pData = NULL) throw(...);
		static void cleanUp() throw(...)
		{
			if (WSACleanup() != 0)
				throw _win32_error(WSAGetLastError());
		}

		static tstring getLocalHostName() throw(...)
		{
			char hostName[200];
			int nameLen = sizeof(hostName);

			THROW_BAD_SOCKET_RET(::gethostname(hostName, nameLen));

			#ifdef _UNICODE
			return asciiToUnicode(hostName);
			#else
			return hostName;
			#endif
		}



	protected:
		void InitConnectEx();
		void clearAddrs()
		{
			memset(&m_localAddr, 0, sizeof(m_localAddr));
			memset(&m_remoteAddr, 0, sizeof(m_remoteAddr));
		}

	// Data
		SOCKET m_socket;
		SockAddrIn m_remoteAddr;	// Addr of the remote host (zeros if listening socket)
		SockAddrIn m_localAddr;		// Addr of this host
		Event m_evtNetwork;

	protected:
		static LPFN_TRANSMITFILE s_transferFilePointer;
	};
}

